package main

import (
	"bufio"
	"fmt"
	"github.com/fatih/color"
	"io"
	"net/http"
	"os"
	"sync"
)

// 全局HTTP客户端
var httpClient = http.Client{
	// 关闭超时，否则会导致下载时间超过超时时间时，TCP连接被自动断开
	Timeout: 0,
	Transport: &http.Transport{
		// 从环境变量读取代理配置
		Proxy: http.ProxyFromEnvironment,
		// 关闭keep-alive确保一个线程就使用一个TCP连接
		DisableKeepAlives: true,
	},
}

// ShardTask 单个分片下载任务的任务参数和状态量
type ShardTask struct {
	// 下载链接
	Url string
	// 分片序号，从1开始
	Order int
	// 下载的文件路径
	FilePath string
	// 分片的起始范围（字节，包含）
	RangeStart int64
	// 分片的结束范围（字节，包含）
	RangeEnd int64
	// 已下载的部分（字节）
	DownloadSize int64
	// 该任务是否完成
	TaskDone bool
}

// NewShardTask 构造函数
func NewShardTask(url string, order int, filePath string, rangeStart int64, rangeEnd int64) *ShardTask {
	return &ShardTask{
		// 设定任务参数
		Url:        url,
		Order:      order,
		FilePath:   filePath,
		RangeStart: rangeStart,
		RangeEnd:   rangeEnd,
		// 初始化状态量
		DownloadSize: 0,
		TaskDone:     false,
	}
}

// DoShardGet 开始下载这个分片（该方法在goroutine中执行）
//
// 将会根据RangeStart和RangeEnd范围，在写入目的文件时，将文件指针Seek到RangeStart位置处开始写入文件
// 实现多个线程同时向一个文件不同位置分别并发写入内容
func (task *ShardTask) DoShardGet(waitGroup *sync.WaitGroup) {
	// 打开文件
	file, e := os.OpenFile(task.FilePath, os.O_WRONLY, 0755)
	if e != nil {
		color.Red("任务%d打开文件失败！", task.Order)
		color.HiRed("%s", e)
		waitGroup.Done()
		return
	}
	// 结束时关闭文件
	defer func() {
		_ = file.Close()
	}()
	// 设定写入偏移量
	_, e = file.Seek(task.RangeStart, io.SeekStart)
	if e != nil {
		color.Red("任务%d设定文件偏移量失败！", task.Order)
		color.HiRed("%s", e)
		waitGroup.Done()
		return
	}
	// 准备请求
	request, e := http.NewRequest("GET", task.Url, nil)
	if e != nil {
		color.Red("任务%d创建请求出错！", task.Order)
		color.HiRed("%s", e)
		waitGroup.Done()
		return
	}
	// 设定请求头
	request.Header.Set("Range", fmt.Sprintf("bytes=%d-%d", task.RangeStart, task.RangeEnd))
	// 发送请求
	response, e := httpClient.Do(request)
	if e != nil {
		color.Red("任务%d发送下载请求出错！", task.Order)
		color.HiRed("%s", e)
		waitGroup.Done()
		return
	}
	// 结束时关闭请求体读取器
	defer func() {
		_ = response.Body.Close()
	}()
	// 读取缓冲区(64KB)
	buffer := make([]byte, 64*1024)
	// 准备写入文件
	writer := bufio.NewWriter(file)
	for {
		// 读取请求体
		// 读取一次内容至缓冲区
		readSize, readError := response.Body.Read(buffer)
		// 处理非EOF错误
		if readError != nil && readError != io.EOF {
			// 如果读取完毕则退出循环
			color.Red("任务%d读取响应错误！", task.Order)
			color.HiRed("%s", readError)
			waitGroup.Done()
			return
		}
		// 读取到内容时，写入文件
		if readSize > 0 {
			// 把缓冲区内容写入文件
			_, writeError := writer.Write(buffer[:readSize])
			if writeError != nil {
				color.Red("任务%d写入文件写入器时出现错误！", task.Order)
				color.HiRed("%s", writeError)
				waitGroup.Done()
				return
			}
			writeError = writer.Flush()
			if writeError != nil {
				color.Red("任务%d写入文件时出现错误！", task.Order)
				color.HiRed("%s", writeError)
				waitGroup.Done()
				return
			}
		}
		// 记录下载进度
		task.DownloadSize += int64(readSize)
		// 若读取到末尾则退出
		if readError == io.EOF {
			break
		}
	}
	// 标记任务完成
	task.TaskDone = true
	// 使线程组中计数器-1
	waitGroup.Done()
}